// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include <nebula/core/buffer.h>
#include <nebula/core/record_batch.h>
#include <turbo/utility/status.h>
#include <nebula/types/type_fwd.h>
#include <turbo/base/macros.h>

namespace nebula {
    template<typename T>
    turbo::Status CopyBufferFromVector(const std::vector<T> &values, MemoryPool *pool,
                                       std::shared_ptr<Buffer> *result) {
        int64_t nbytes = static_cast<int>(values.size()) * sizeof(T);

        TURBO_MOVE_OR_RAISE(auto buffer, allocate_buffer(nbytes, pool));
        auto immutable_data = reinterpret_cast<const uint8_t *>(values.data());
        std::copy(immutable_data, immutable_data + nbytes, buffer->mutable_data());
        memset(buffer->mutable_data() + nbytes, 0,
               static_cast<size_t>(buffer->capacity() - nbytes));

        *result = std::move(buffer);
        return turbo::OkStatus();
    }

    // Sets approximately pct_null of the first n bytes in null_bytes to zero
    // and the rest to non-zero (true) values.
    TURBO_EXPORT void random_null_bytes(int64_t n, double pct_null,
                                                 uint8_t *null_bytes);

    TURBO_EXPORT void random_is_valid(int64_t n, double pct_null,
                                               std::vector<bool> *is_valid,
                                               int random_seed = 0);

    TURBO_EXPORT void random_bytes(int64_t n, uint32_t seed, uint8_t *out);

    TURBO_EXPORT std::string random_string(int64_t n, uint32_t seed);

    TURBO_EXPORT int32_t DecimalSize(int32_t precision);

    TURBO_EXPORT void random_ascii(int64_t n, uint32_t seed, uint8_t *out);

    TURBO_EXPORT int64_t CountNulls(const std::vector<uint8_t> &valid_bytes);

    TURBO_EXPORT turbo::Status MakeRandomByteBuffer(int64_t length, MemoryPool *pool,
                                                             std::shared_ptr<ResizableBuffer> *out,
                                                             uint32_t seed = 0);

    TURBO_EXPORT uint64_t random_seed();

#define DECL_T() typedef typename TestFixture::T T;

#define DECL_TYPE() typedef typename TestFixture::Type Type;

    // ----------------------------------------------------------------------
    // A RecordBatchReader for serving a sequence of in-memory record batches

    class BatchIterator : public RecordBatchReader {
    public:
        BatchIterator(const std::shared_ptr<Schema> &schema,
                      const std::vector<std::shared_ptr<RecordBatch> > &batches)
            : schema_(schema), batches_(batches), position_(0) {
        }

        std::shared_ptr<Schema> schema() const override { return schema_; }

        turbo::Status read_next(std::shared_ptr<RecordBatch> *out) override {
            if (position_ >= batches_.size()) {
                *out = nullptr;
            } else {
                *out = batches_[position_++];
            }
            return turbo::OkStatus();
        }

    private:
        std::shared_ptr<Schema> schema_;
        std::vector<std::shared_ptr<RecordBatch> > batches_;
        size_t position_;
    };

    static inline std::vector<std::shared_ptr<DataType> (*)(FieldVector, std::vector<int8_t>)>
    UnionTypeFactories() {
        return {sparse_union, dense_union};
    }

    // Return the value of the NEBULA_TEST_DATA environment variable or return error
    // turbo::Status
    TURBO_EXPORT turbo::Status GetTestResourceRoot(std::string *);

    // Return the value of the NEBULA_TIMEZONE_DATABASE environment variable
    TURBO_EXPORT std::optional<std::string> GetTestTimezoneDatabaseRoot();

    // Set the Timezone database based on the NEBULA_TIMEZONE_DATABASE env variable
    // This is only relevant on Windows, since other OSs have compatible databases built-in
    TURBO_EXPORT turbo::Status InitTestTimezoneDatabase();

    // Get a TCP port number to listen on.  This is a different number every time,
    // as reusing the same port across tests can produce spurious bind errors on
    // Windows.
    TURBO_EXPORT int GetListenPort();

    // Get a IPv4 "address:port" to listen on.  The address will be a loopback
    // address.  Compared to GetListenPort(), this will minimize the risk of
    // port conflicts.
    TURBO_EXPORT std::string GetListenAddress();

    TURBO_EXPORT
    const std::vector<std::shared_ptr<DataType> > &all_dictionary_index_types();

    // Get a list of supported hardware flags from the given candidates.
    // The result will always contain 0, meaning no optional CPU feature enabled at all.
    TURBO_EXPORT
    std::vector<int64_t> GetSupportedHardwareFlags(
        const std::vector<int64_t> &candidate_flags);
} // namespace nebula
